import * as tfvis from "@tensorflow/tfjs-vis";
import * as tf from "@tensorflow/tfjs";
import { getData } from "./data";

(async () => {
  let data = getData(400);
  tfvis.render.scatterplot(
    {
      name: "xor数据",
    },
    {
      values: [
        data.filter((p) => p.label === 1),
        data.filter((p) => p.label === 0),
      ],
    }
  );

  let model = tf.sequential();

  // 隐藏层 这层的输出是下一层输入
  model.add(
    tf.layers.dense({
      units: 4, // 4个神经元
      inputShape: [2],
      activation: "relu", // 激活函数，不加这个线性组合叠加还是线性，这样不能解决线性问题
    })
  );
  // 输出层
  model.add(
    tf.layers.dense({
      units: 1,
      activation: "sigmoid",
    })
  );

  model.compile({
    loss: tf.losses.logLoss,
    optimizer: tf.train.adam(0.1),
  });

  let inputs = tf.tensor(data.map((p) => [p.x, p.y]));
  let labels = tf.tensor(data.map((p) => p.label));

  // 获取表单的提交
  window["predict"] = () => {
    let formData = document.form;
    let pred = model.predict(
      tf.tensor([[formData.x.value * 1, formData.y.value * 1]])
    ); // [ [] ]
    console.log(pred.dataSync()[0]);
    return false;
  };

  await model.fit(inputs, labels, {
    batchSize: 40,
    epochs: 500,
    callbacks: tfvis.show.fitCallbacks(
      {
        name: "训练结果",
      },
      ["loss"]
    ),
  });
  console.log("训练完成");
})();
